"""
图像处理模块 - 处理图像上传和识别
"""

import streamlit as st
import os
import hashlib
import base64
from datetime import datetime


from modules.config import UPLOADED_IMAGES_DIR

@st.cache_resource
def load_image_classifier():
    """加载图像分类模型（使用TensorFlow Lite）"""
    try:
        import numpy as np
        from PIL import Image
        
        class TFLiteClassifier:
            def __init__(self):
                print("初始化TensorFlow Lite文物分类器模拟")
                
                # 文物类别
                self.artifacts = ["狮子", "老虎", "大象", "长颈鹿", "斑马", "企鹅", "北极熊", "海豚", "孔雀", "金丝猴"]
                
                # 动物特征（简化的HSV颜色特征）
                self.artifact_features = {
                    "狮子": [25, 35, 45],  # 金黄色毛发，中等饱和度
                    "老虎": [20, 30, 50],  # 橙黄色条纹，较高饱和度
                    "大象": [35, 25, 60],  # 灰色皮肤，低饱和度
                    "长颈鹿": [45, 40, 55],  # 黄褐色斑点，中等饱和度
                    "斑马": [30, 20, 70],  # 黑白条纹，高对比度
                    "企鹅": [20, 15, 80],  # 黑白羽毛，高对比度
                    "北极熊": [25, 10, 85],  # 白色毛发，极高亮度
                    "海豚": [40, 30, 65],  # 深灰色皮肤，中等亮度
                    "孔雀": [50, 60, 50],  # 蓝绿色羽毛，高饱和度
                    "金丝猴": [35, 45, 55]   # 金黄色毛发，中等饱和度
                }
            
            def analyze_image(self, image_path):
                try:
                    # 加载并处理图像
                    with Image.open(image_path) as img:
                        # 调整大小
                        img = img.resize((224, 224))
                        # 转换为HSV色彩空间（模拟）
                        img_rgb = np.array(img.convert('RGB'))
                        # 提取RGB通道平均值作为简单特征
                        r_mean = np.mean(img_rgb[:,:,0])
                        g_mean = np.mean(img_rgb[:,:,1])
                        b_mean = np.mean(img_rgb[:,:,2])
                        
                        # 简单的HSV转换模拟（这不是准确的HSV转换，仅为示例）
                        h_sim = (r_mean + g_mean) / 2
                        s_sim = (g_mean + b_mean) / 2
                        v_sim = (r_mean + b_mean) / 2
                        img_feature = [h_sim, s_sim, v_sim]
                        
                        # 计算与各文物特征的距离
                        distances = {}
                        for artifact, feature in self.artifact_features.items():
                            dist = np.sqrt(sum((a-b)**2 for a, b in zip(img_feature, feature)))
                            distances[artifact] = dist
                        
                        # 寻找最接近的文物
                        closest_artifact = min(distances.items(), key=lambda x: x[1])[0]
                        
                        # 计算置信度（距离的倒数，归一化）
                        min_dist = distances[closest_artifact]
                        max_dist = max(distances.values())
                        confidence = 1.0 - (min_dist / (max_dist + 1e-5))
                        # 限制在0.6-0.95之间，增加随机性
                        confidence = 0.6 + 0.35 * confidence + np.random.normal(0, 0.05)
                        confidence = min(0.95, max(0.6, confidence))
                        
                        return {"result": closest_artifact, "confidence": float(confidence)}
                except Exception as e:
                    print(f"图像分析出错: {str(e)}")
                    return {"result": "未知动物", "confidence": 0.1}
        
        classifier = TFLiteClassifier()
        print("TFLite文物分类器模拟加载成功")
        return classifier
    except Exception as e:
        st.error(f"⚠️ 加载图像分类模型失败: {str(e)}")
        return None

def save_uploaded_image(uploaded_file):
    """保存上传的图片并返回文件路径"""
    if uploaded_file is None:
        return None
    
    try:
        file_size = len(uploaded_file.getvalue())
        max_size = 10 * 1024 * 1024  # 10MB
        
        if file_size > max_size:
            st.error(f"文件过大！请上传小于10MB的图片。当前文件大小: {file_size/1024/1024:.2f}MB")
            return None
            
        file_extension = os.path.splitext(uploaded_file.name)[1]
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        unique_id = hashlib.md5((timestamp + uploaded_file.name).encode()).hexdigest()
        filename = f"{unique_id}{file_extension}"
        filepath = os.path.join(UPLOADED_IMAGES_DIR, filename)
        
        with open(filepath, "wb") as f:
            f.write(uploaded_file.getbuffer())
        
        return filepath
    except Exception as e:
        st.error(f"⚠️ 保存图像过程出错: {str(e)}")
        return None

def classify_artifact_with_ernie(client, image_path, kg_data):
    """使用ERNIE VL模型进行文物图像识别"""
    if client is None or kg_data is None:
        return "未知动物", 0.0
    
    try:
        with open(image_path, "rb") as image_file:
            base64_image = base64.b64encode(image_file.read()).decode('utf-8')
        
        prompt = """你是一位专业的动物识别专家。请分析这张图片中的动物，并告诉我：
                    1. 这是什么动物？（如：狮子、老虎、大象等）
                    2. 你的识别置信度（0.0-1.0之间的数值）
                    只需简单回答：动物名称，置信度。例如：
                    狮子，0.95
                """
        
        messages = [
            {
                'role': 'user', 
                'content': [
                    {"type": "text", "text": prompt},
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}
                    }
                ]
            }
        ]
        
        completion = client.chat.completions.create(
            model="ernie-4.5-turbo-vl-32k",
            messages=messages,
            temperature=0.1,
        )
        
        response = completion.choices[0].message.content
        
        artifact_name = "未知动物"
        confidence = 0.5
        
        try:
            response = response.strip().replace('，', ',').replace('：', ':')
            
            if ',' in response:
                parts = response.split(',', 1)
                artifact_name = parts[0].strip()
                import re
                conf_matches = re.findall(r'0\.\d+', parts[1])
                if conf_matches:
                    confidence = float(conf_matches[0])
            
            if artifact_name != "未知动物" and confidence == 0.5:
                confidence = 0.8
                
            artifact_names = [node['name'] for node in kg_data['nodes'] if node['type'] == 'animal']
            
            exact_match = False
            for known_artifact in artifact_names:
                if artifact_name == known_artifact:
                    exact_match = True
                    break
            
            if not exact_match:
                best_match = None
                best_score = 0
                
                for known_artifact in artifact_names:
                    match_count = 0
                    for char in artifact_name:
                        if char in known_artifact:
                            match_count += 1
                    
                    score = match_count / max(len(artifact_name), len(known_artifact))
                    
                    if score > best_score and score > 0.5: 
                        best_score = score
                        best_match = known_artifact
                
                if best_match:
                    artifact_name = best_match
                    confidence = max(0.6, confidence * 0.9)
        
        except Exception as e:
            st.error(f"解析ERNIE响应时出错: {str(e)}")
            if "狮子" in response:
                artifact_name = "狮子"
                confidence = 0.7
            elif "老虎" in response:
                artifact_name = "老虎"
                confidence = 0.7
            elif "大象" in response:
                artifact_name = "大象"
                confidence = 0.7
        
        return artifact_name, confidence
        
    except Exception as e:
        st.error(f"⚠️ ERNIE VL文物识别过程出错: {str(e)}")
        return "未知文物", 0.0

def img_to_base64(image_path):
    """将图片转换为Base64编码"""
    try:
        with open(image_path, "rb") as img_file:
            return base64.b64encode(img_file.read()).decode('utf-8')
    except Exception:
        return None

@st.cache_data
def get_artifact_sample_images():
    """获取示例动物图片（本地路径）"""
    artifact_images = {
        "狮子": "https://img95.699pic.com/photo/60034/6765.jpg_wh860.jpg",
        "老虎": "https://ts1.tc.mm.bing.net/th/id/R-C.52c5977353a7674c6871f746375b2b37?rik=3qvKbklVf7bYDA&riu=http%3a%2f%2fwww.v3wall.com%2fwallpaper%2f1680_1050%2f1002%2f1680_1050_20100208121107627357.jpg&ehk=9mLa%2bzcgzZizA%2fmOADkC%2bPRXaleZtzV590JxIsG68ls%3d&risl=&pid=ImgRaw&r=0",
        "大象": "https://so1.360tres.com/t01e4070de9d4128739.jpg",
        "长颈鹿": "https://ts1.tc.mm.bing.net/th/id/OIP-C._pwRGiFOLoFkYYZV9ysfcAHaEp?r=0&rs=1&pid=ImgDetMain&o=7&rm=3",
        "斑马": "https://n.sinaimg.cn/sinacn21/214/w2048h1366/20180730/8950-hhacrcc7490948.jpg",
        "企鹅": "https://ts2.tc.mm.bing.net/th/id/OIP-C.__HKN_B8Do9eUO96F3tkQgHaE7?r=0&rs=1&pid=ImgDetMain&o=7&rm=3",
    }
    return artifact_images 